package optimizer;

import midcode.MidCode;
import midcode.MidcodeGenerator;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.Map;

public class RegPool {
    private final HashMap<Integer, String> SReg;
    private HashMap<Integer, String> TReg;
    private final int calTNum = 3;
    private final HashMap<String, ArrayList<Integer>> tmpMap;

    public RegPool() {
        SReg = new HashMap<>();
        for (int i = 0; i < 8; i++) {
            SReg.put(i, "");
        }
        TReg = new HashMap<>();
        for (int i = 0; i < 10; i++) {
            TReg.put(i, "");
        }
        tmpMap = MidcodeGenerator.getInstance().getTmpMap();
    }


    public String allocSReg(String var) {
        for (Map.Entry<Integer, String> entry : SReg.entrySet()) {
            if (entry.getValue().equals("")) {
                entry.setValue(var);
                return "$s" + entry.getKey();
            }
        }
        return "";
    }

    public String allocTReg(String var, int index) {
        if (!Optimizer.isOptimise || !var.startsWith("T_")) {
            return "$t0";
        }
        for (Map.Entry<Integer, String> entry : TReg.entrySet()) {
            if (entry.getKey() < calTNum) {
                continue;
            }
            if (entry.getValue().equals("")) {
                entry.setValue(var);
                return "$t" + entry.getKey();
            } else {
                boolean switchFlag = true;
                for (Integer i : tmpMap.get(entry.getValue())) {
                    if (index <= i) {
                        switchFlag = false;
                        break;
                    }
                }
                if (switchFlag) {
                    entry.setValue(var);
                    return "$t" + entry.getKey();
                }
            }
        }
        return "$t0";
    }

    public String getTReg(String var) {
        for (Map.Entry<Integer,String> entry : TReg.entrySet()) {
            if (entry.getValue().equals(var)) {
                return "$t" + entry.getKey();
            }
        }
        return "$t0";
    }

    public ArrayList<String> writeBackTReg(int index) {
        ArrayList<String> list = new ArrayList<>();
        for (Map.Entry<Integer,String> entry : TReg.entrySet()) {
            if (entry.getKey() < calTNum) {
                continue;
            }
            if (!entry.getValue().equals("")) {
                boolean writeBack = false;
                for (Integer i : tmpMap.get(entry.getValue())) {
                    if (index <= i) {
                        writeBack = true;
                        break;
                    }
                }
                if (writeBack) {
                    list.add(entry.getValue());
                    list.add("$t" + entry.getKey());
                }
            }
        }
        return list;
    }

    public void clearTReg() {
        TReg = new HashMap<>();
        for (int i = 0; i < 10; i++) {
            TReg.put(i, "");
        }
    }
}

